//
// Created by hjhu on 2024/7/2.
//

#include "ConvBlock.h"

namespace coastal {

    ConvBlockImpl::ConvBlockImpl(int64_t
                         in_planes,
                         int64_t out_planes
    )
            :
            bn1(in_planes),
            conv1(conv3x3(in_planes, out_planes / 2)),
            bn2(out_planes / 2),
            conv2(conv3x3(out_planes / 2, out_planes / 4)),
            bn3(out_planes / 4),
            conv3(conv3x3(out_planes / 4, out_planes / 4)) {

        if (in_planes != out_planes) {
            downsample = nn::Sequential(
                    nn::BatchNorm2d(in_planes),
                    nn::ReLU(true),
                    nn::Conv2d(nn::Conv2dOptions(in_planes, out_planes, 1).stride(1).bias(false))
            );
        } else {
            downsample = nullptr;
        }

        register_module("bn1", bn1);
        register_module("conv1", conv1);
        register_module("bn2", bn2);
        register_module("conv2", conv2);
        register_module("bn3", bn3);
        register_module("conv3", conv3);
        if (downsample.has_value()) {
            register_module("downsample", downsample.value());
        }
    }
}